/**
 * @program: LoveDataStructures1
 * @description: [用一句话描述此类]
 * @author: wxc
 * @create: 2021-07-27 09:29
 **/
public class AVLTree<E> extends BalanceBinarySearchTree<E> {
    public AVLTree(IComparator<E> comparator) {
        this.comparator = comparator;
    }

    public AVLTree() {
    }

    @Override
    protected Node<E> createNode(E element, Node<E> left, Node<E> right, Node<E> parent) {
        return new AVLNode<>(element, left, right, parent);
    }

    protected static class AVLNode<E> extends Node<E> {
        // 叶子结点默认1
        int height = 1;

        public AVLNode(E element, Node<E> left, Node<E> right, Node<E> parent) {
            super(element, left, right, parent);
        }

        public int balanceFactor() {
            // O(1)
            int leftHeight = this.left == null ? 0 : ((AVLNode<E>) this.left).height;
            int rightHeight = this.right == null ? 0 : ((AVLNode<E>) this.right).height;
            return leftHeight - rightHeight;
        }

        public void updateHeight() {
            // O(1)
            int leftHeight = this.left == null ? 0 : ((AVLNode<E>) this.left).height;
            int rightHeight = this.right == null ? 0 : ((AVLNode<E>) this.right).height;
            this.height = 1 + Math.max(leftHeight, rightHeight);
        }

        // 找到高度更高的那个子结点
        public Node<E> getTallerChildNode() {
            int leftHeight = this.left == null ? 0 : ((AVLNode<E>) this.left).height;
            int rightHeight = this.right == null ? 0 : ((AVLNode<E>) this.right).height;
            if (leftHeight < rightHeight) return this.right;
            if (leftHeight > rightHeight) return this.left;
                // 我觉得reBalance不会进下面的else
            else {
                if (this.isLeftChildNode()) return this.left;
                return right;
            }
        }
    }

    // 该结点是否平衡
    private boolean isBalanced(Node<E> node) {
        return Math.abs(((AVLNode<E>) node).balanceFactor()) <= 1;
    }

    // 更新高度
    private void updateHeight(Node<E> node) {
        ((AVLNode<E>) node).updateHeight();
    }

    // 恢复平衡
    // grandNode是高度最低的那个不平衡的那个结点
    private void reBalance(Node<E> grandNode) {
        // 找到parentNode
        Node<E> parentNode = ((AVLNode<E>) grandNode).getTallerChildNode();
        Node<E> node = ((AVLNode<E>) parentNode).getTallerChildNode();
        // L
        if (grandNode.left == parentNode) {
            // LL
                if (parentNode.left == node) {
                rotateRight(grandNode);
            }
            // LR
            else {
                rotateLeft(parentNode);
                rotateRight(grandNode);
            }
        }
        // R
        else {
            // RR
            if (parentNode.right == node) {
                rotateLeft(grandNode);
            }
            // RL
            else {
                rotateRight(parentNode);
                rotateLeft(grandNode);
            }
        }
    }

    @Override
    protected void afterRotate(Node<E> childRightChildNode, Node<E> childNode, Node<E> node) {
        super.afterRotate(childRightChildNode, childNode, node);
        // 更新高度 不能换位置
        updateHeight(node);
        updateHeight(childNode);
    }

    @Override
    protected void afterAdd(Node<E> node) {
        // node一定是叶子结点
        // node不可能为null
        // 找到最先失衡结点
        while (node.parent != null) {
            node = node.parent;
            // 如果结点平衡 则更新高度
            if (isBalanced(node)) updateHeight(node);
                // 如果node结点失衡则恢复平衡并break
            else {
                reBalance(node);
                break;
            }
        }
    }

    @Override
    protected void afterRemove(Node<E> node,Node<E> replaceNode) {
        while (node.parent != null) {
            node = node.parent;
            // 如果结点平衡 则更新高度
            if (isBalanced(node)) updateHeight(node);
                // 如果node结点失衡则恢复平衡
            else {
                reBalance(node);
            }
        }
    }

    @Override
    public String string(Node<E> node) {
        return super.string(node) + "h(" + ((AVLNode<E>)node).height + ")";
    }
}
